import re
import json
import random
from openai import OpenAI
from tqdm import tqdm

# 设置API
openai_api_key = "EMPTY"
openai_api_base = "http://localhost:8000/v1"
client = OpenAI(api_key=openai_api_key, base_url=openai_api_base)


# 调用本地模型
def qwen3_8B_local(topic, if_think=False):
    # 替换用户提示中的变量
    user_msg = user_msg_template.format(TOPIC=topic)

    # 构造 message 列表（包含 system 和 user 两种角色）
    messages = [
        {"role": "system", "content": system_msg},
        {"role": "user", "content": user_msg}
    ]
    temperature = round(random.uniform(0.6, 1.4), 2)
    top_p = round(random.uniform(0.85, 0.99), 2)
    # 调用本地 API
    res = client.chat.completions.create(
        model="/home/work_nfs11/asr_data/ckpt/Qwen3-8B",
        messages=messages,
        max_tokens=32768,
        temperature=temperature,
        top_p=top_p,
        extra_body={
            "top_k": 20,
            "chat_template_kwargs": {"enable_thinking": if_think},
        },
    ).choices[0].message.content

    # 去除模型输出中的 <think> 标签（如果有）
    # res = re.sub(r'<think>.*?</think>', '', res, flags=re.DOTALL).strip()
    return res


# system prompt
system_msg = '''你是一个人类对话数据专家，擅长模仿真实的人与人之间的口语对话。
说话风格要非常自然，并符合上下文情境。
重要提示：考虑一种情境，在历史对话之后有一个当前回合的中性情感文本，但可以有不同可能的说话风格。不同的当前说话风格会使回应在语义上产生较大差异。每个回合只说一句话，句子应是口语化的、自然流畅的，不要太正式。
[必须遵守的规则]：
0. 我们使用特殊标记 <> 来表示你需要生成的类别类型。输出中不要包含 <>。
1. 在对话中使用多样化的说话风格。
2. 当前回合的文本情感为中性，回应需要根据当前回合的内容自然地做出反应，不能只是简单复制当前回合的风格。
3. 对话中有两个说话人 A 和 B，他们交替发言。
4. 每个回合格式为：<说话人> (<语气>): <文本>
5. 回合顺序为：历史回合 -> 当前回合 -> 接下来的回应。
6. 对话的过渡要非常连贯，符合常识。
7. 对话中要有情绪变化。
8. 输出的有效字典格式如下：
{
  "current_turn_style_1": "<说话人> (<语气>>): <文本>",
  "response_of_current_style_1": "<说话人> (<语气>): <文本>",
}
9. 输出一个有效的字典示例，以便可以被解析为字典。
10. 说话人只能使用 A 或 B。说话人 A 是"current_turn_style"，而说话人 B 是"response_of_current_style"。说话人 B 的文本长度是说话人 A 的三倍以上。
11. 你只能使用以下定义的说话风格进行表示（<语气>）。非常重要！不要使用其他未定义的类别！！！
  - 11.1 语气：中性、生气、高兴、伤心、害怕、惊讶、厌恶、疑惑、嘲讽、尴尬、好奇、担忧、害羞、抱歉。
  - 11.2 说话人：A 或 B。
12.请注意，不要总是使用 (<中性>)。请尝试使用多种不同语气。
'''

# user prompt模板
user_msg_template = '''
[topic]：{TOPIC}
请基于主题[topic]生成一轮带有说话相关情绪风格的对话。请预测接下来的回应。
对话主题为[topic]。您可以自由想象对话内容，但应基于常识。我们使用(<语气>)来表示说话风格。
'''

# 填充主题与历史轮数
# 可用话题
topics = [
    "智能家居", "旅游计划", "天气变化", "运动健身", "教育理念", "感情关系", "亲子教育",
    "职场压力", "节日安排", "看病就医", "健康饮食", "娱乐明星", "网络购物", "环保出行",
    "手机换新", "宠物饲养", "学习技巧", "租房经历", "理财方式", "朋友聚会", "婚姻观念",
    "电影推荐", "音乐品味", "AI 技术", "未来职业", "健身习惯", "科技产品", "智能助手",
    "社交礼仪", "读书分享", "心理健康", "家装风格", "孩子教育", "上班通勤", "饮食误区"
]
more_topics = [
    "职业选择", "大学生活", "网恋经历", "考研准备", "异地恋", "家庭聚会", "室友相处", "早起困难", "时间管理",
    "二手交易", "校园食堂", "网课体验", "实习经历", "跳槽决定", "搬家烦恼", "交通工具选择", "社恐日常", "表白经历",
    "办公室政治", "请假理由", "追星日常", "手机依赖", "节食减肥", "摄影爱好", "咖啡选择", "夜生活", "考证经验",
    "护肤心得", "穿搭建议", "消费主义", "书籍借阅", "网络诈骗", "打游戏", "直播打赏", "租房被骗", "房价讨论",
    "装修预算", "邻里纠纷", "独居生活", "中餐烹饪", "西餐餐桌礼仪", "假期安排", "跨年计划", "家庭教育理念",
    "儿童电子产品", "社交媒体焦虑", "人工智能影响", "远程办公", "元宇宙", "智能翻译", "智能手表", "语音助手对比",
    "电动汽车", "自动驾驶", "新能源汽车补贴", "短视频沉迷", "碎片化阅读", "知识付费", "虚拟偶像", "情绪管理",
    "心理咨询", "睡眠障碍", "压力释放", "数字健康", "智能医疗", "饮料选择", "茶文化", "节能环保", "垃圾分类",
    "地摊经济", "创业失败", "副业选择", "个人品牌打造", "网络求职", "简历优化", "面试技巧", "职场沟通",
    "开会文化", "绩效考核", "升职加薪", "办公软件推荐", "笔记技巧", "时间块管理", "专注力提升", "自律习惯",
    "阅读打卡", "英语学习", "口语练习", "留学申请", "签证流程", "文化差异", "旅游攻略", "旅行拍照", "打包清单",
    "护照办理", "出境购物", "航班延误", "海关安检", "城市印象", "旅途中遇到的趣事", "酒店选择", "打车体验"
]

more_topics_50 = [
    "假期旅行计划", "追剧推荐", "同学聚会", "新冠后遗症", "节能家电", "家庭矛盾", "表弟高考", "宠物走失",
    "快递丢失", "吃外卖健康吗", "网络成瘾", "学习拖延", "手机维修", "笔记本推荐", "运动手环", "戴口罩习惯",
    "疫情改变生活", "公积金提取", "租房押金纠纷", "校园霸凌", "长辈微信使用", "情侣吵架", "分手挽回",
    "约会地点选择", "恋爱开销", "婚礼准备", "彩礼观念", "失恋恢复", "单身焦虑", "母婴用品选购",
    "孩子沉迷游戏", "老年人看病", "老年人智能手机使用", "乡村振兴", "生育政策", "晚婚晚育", "非遗文化",
    "汉服穿搭", "跨文化交流", "中英文混搭", "写作能力提升", "公开演讲恐惧", "视频剪辑学习", "短视频创作",
    "街头采访", "电商平台对比", "会员制陷阱", "低头族", "盲盒消费", "校园恋爱"
]

more_topics_50_extra = [
    "居家办公体验", "买房还是租房", "相亲经历", "加班文化", "公司团建", "年度体检", "医保报销流程",
    "公园遛狗", "共享单车体验", "校园兼职", "大学选专业", "学区房焦虑", "城市与乡村生活对比",
    "双十一购物", "家庭收纳技巧", "情绪价值", "脱发焦虑", "熬夜的危害", "新媒体运营", "流量明星争议",
    "人工智能替代就业", "AI 写作工具", "智能语音助手体验", "大模型生成内容", "网络暴力", "网络隐私保护",
    "人脸识别技术", "指纹支付", "智能门锁使用", "手机续航焦虑", "数字人民币", "区块链应用", "元宇宙社交",
    "虚拟现实设备", "智能健身镜", "智能马桶", "高考志愿填报", "博士毕业就业", "中年危机", "退休规划",
    "父母情绪管理", "婚姻冷静期", "互联网脱口秀", "短剧内容爆火", "AI 换脸技术", "网络谣言传播",
    "社交账号管理", "电子阅读 vs 纸质书", "网络语言文化", "校园食堂吐槽"
]

more_topics_100 = [
    "考驾照经验", "通勤路上的烦恼", "早餐吃什么", "晚餐点外卖还是做饭", "如何节约用电", "亲戚催婚怎么办",
    "银行办卡流程", "信用卡还款", "交通违章处理", "城市限号", "理发店踩坑", "夜市小吃推荐", "网红餐厅排队",
    "奶茶饮品选择", "超市囤货", "月度账单复盘", "双语教育利弊", "幼儿园入园焦虑", "高考复读压力", "中年转行",
    "职业倦怠", "裸辞后的生活", "996 工作制", "远程办公工具", "工位风水", "公司食堂体验", "办公区域冲突",
    "职场内耗", "职场PUA", "领导画饼", "上班摸鱼技巧", "工作日报怎么写", "线上会议效率", "文件命名习惯",
    "求职被骗经历", "离职后的社保", "带薪年假被拒", "组内晋升机制", "午休要不要睡", "商务礼仪误区",
    "电子产品过度使用", "微信群礼仪", "朋友圈屏蔽谁", "微博热搜真假", "视频号与抖音", "直播打工人",
    "网络追剧与补番", "漫画推荐", "剧本杀初体验", "密室逃脱好玩吗", "桌游社交", "KTV 歌单推荐",
    "考研住宿问题", "留学生归国适应", "学术造假新闻", "师生关系边界", "论文查重焦虑", "期末考复习计划",
    "图书馆占座大战", "高校食堂大比拼", "实验室的人际关系", "课程评价机制", "毕业季失落", "同学间的借钱纠纷",
    "租房合同细节", "搬家时丢东西", "楼道噪音问题", "邻居孩子吵闹", "小区物业收费", "快递放门口安全吗",
    "宠物邻里矛盾", "外卖员送错单", "停车位争执", "电梯里遇见前任", "便利店深夜营业", "打车被绕路",
    "滴滴与网约车体验", "地铁上的尴尬", "共享充电宝价格", "旅游景点踩坑", "导游强制购物", "跟团游和自由行",
    "旅行拍照修图", "语言不通怎么办", "护照丢了怎么处理", "住宿退订纠纷", "酒店卫生问题", "航班取消应对",
    "旅行险购买", "旅途结识新朋友", "异国风俗冲突", "旅游纪念品挑选", "景区厕所难题", "登山装备准备",
    "露营注意事项", "出行天气影响", "旅行App推荐", "旅游直播",
]
more_topics_100_extra = [
    "买房经历", "二手车交易", "婚礼策划", "健身计划", "宠物训练", "孩子早教", "家庭理财", "老人赡养",
    "网络安全", "个人隐私保护", "线上学习体验", "工作压力释放", "同事关系处理", "办公室恋情", "升职加薪",
    "职场竞争", "自我提升方法", "兴趣爱好培养", "旅行攻略", "摄影技巧", "美食制作", "家庭装修",
    "租房合同", "车险选择", "节假日出游", "气候变化", "环保生活", "新能源车", "数字货币",
    "社交媒体影响", "网络成瘾", "游戏文化", "电竞比赛", "影视剧推荐", "读书笔记", "写作技巧",
    "时间管理", "睡眠质量", "健康饮食", "心理咨询", "情感困惑", "婚姻生活", "亲子沟通",
    "人际交往", "公共交通", "疫情防控", "医疗保险", "保险理赔", "法律咨询", "交通事故处理",
    "消费陷阱", "理财产品", "股票投资", "基金购买", "养老规划", "子女教育", "留学申请",
    "语言学习", "文化差异", "移民生活", "海外就业", "跨境电商", "数字营销", "人工智能应用",
    "自动化技术", "智能家居", "云计算", "大数据分析", "网络购物体验", "快递服务", "物流配送",
    "节能减排", "社区活动", "志愿服务", "艺术欣赏", "音乐学习", "体育赛事", "户外运动",
    "极限运动", "电影制作", "戏剧表演", "手工艺品", "科技新品发布", "新车评测", "健康养生",
    "减肥经验", "饮食误区", "心理健康", "职业规划", "创业故事", "团队合作", "领导力培养"
]

all_topics = topics + more_topics + more_topics_50 + more_topics_50_extra + more_topics_100 + more_topics_100_extra
print(f"总话题数量：{len(all_topics)}")

# 目标样本数
# target_num = 10000
# output_path = "dialog_data_1000.jsonl"

# with open(output_path, "w", encoding="utf-8") as fout:
#     for i in tqdm(range(target_num), desc="生成对话样本"):
#         topic = random.choice(all_topics)
#         try:
#             response_text = qwen3_8B_local(topic)
#             data = json.loads(response_text)
#             json.dump(data, fout, ensure_ascii=False)
#             fout.write("\n")
#         except Exception as e:
#             print(f"⚠️ 第 {i+1} 条出错：{e}")
#             continue

# print(f"\n✅ 已成功生成 {target_num} 条对话数据，保存在：{output_path}")

target_num = 50000
output_path = "dialog_data_50000.jsonl"
generated_count = 0
topic_index = 0

# with open(output_path, "w", encoding="utf-8") as fout:
#     pbar = tqdm(total=target_num, desc="生成对话样本")
#
#     while generated_count < target_num:
#         topic = all_topics[topic_index % len(all_topics)]  # 循环遍历话题
#         topic_index += 1
#
#         try:
#             response_text = qwen3_8B_local(topic)
#             data = json.loads(response_text)
#             json.dump(data, fout, ensure_ascii=False)
#             fout.write("\n")
#             generated_count += 1
#             pbar.update(1)
#         except Exception as e:
#             print(f"⚠️ 第 {generated_count + 1} 条出错（话题: {topic}）：{e}")
#             continue

# pbar.close()
# print(f"\n✅ 已成功生成 {target_num} 条对话数据，保存在：{output_path}")
# response_text = qwen3_8B_local("心理健康", if_think=True)
# print(response_text)
# response_text = qwen3_8B_local("心理健康", if_think=False)
# print(response_text)


def common_chat_qwen3_8B(user_prompt, system_prompt="you are QWEN3-8B",if_think=False):
    # 构造 message 列表（包含 system 和 user 两种角色）
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt}
    ]
    temperature = round(random.uniform(0.6, 1.4), 2)
    top_p = round(random.uniform(0.85, 0.99), 2)
    # 调用本地 API
    res = client.chat.completions.create(
        model="/home/work_nfs11/asr_data/ckpt/Qwen3-8B",
        messages=messages,
        max_tokens=32768,
        temperature=temperature,
        top_p=top_p,
        extra_body={
            "top_k": 20,
            "chat_template_kwargs": {"enable_thinking": if_think},
        },
    ).choices[0].message.content

    # 去除模型输出中的 <think> 标签（如果有）
    # res = re.sub(r'<think>.*?</think>', '', res, flags=re.DOTALL).strip()
    return res


while True:
    input_text = input("请输入对话内容：")
    if input_text == "exit":
        break
    response_text = common_chat_qwen3_8B(input_text, if_think=False)
    print(f"QWEN3-8B 机器人：{response_text}")

# 使用尽可能简单的对话进行回答，尽可能参考副语言信息进行回答，比如对老人和小孩进行适当的回复，对情感问题进行情绪化回复，对声音事件进行适当的回复，对女性或者男性给出适当的回答等。<think>用户说的话是:这台新买的智能冰箱该怎么设置才能最方便使用？,年龄为:<OLD>,性别为:<MALE>,风格为:<日常口语>,情感为:<NEUTRAL>,声音事件为:<OTHER>,推测使用的回复情感为:<HAPPY>,我应该综合用户的语义和副语言信息给出专业且对应的回答<think end>

# 用户输入内容：这台新买的智能冰箱该怎么设置才能最方便使用？用户年龄: 老人，性别: 女性，风格: 日常口语，情感: 无表情，声音事件: 无声音，推测使用的回复情感: 无表情，请给出回答：